package cc.yihy.spider;

import org.springframework.beans.BeansException;
import org.springframework.boot.actuate.endpoint.PublicMetrics;
import org.springframework.boot.actuate.metrics.Metric;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import us.codecraft.webmagic.Spider;
import us.codecraft.webmagic.monitor.SpiderMonitor;
import us.codecraft.webmagic.scheduler.MonitorableScheduler;

import java.util.Collection;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

public class SpiderPublicMetrics implements ApplicationContextAware, PublicMetrics {

    ApplicationContext context;
    private final static String SPIDER_SUFFIX = "spider.";


    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.context = applicationContext;
    }

    @Override
    public Collection<Metric<?>> metrics() {
        Set<Metric<?>> metricSet = new LinkedHashSet<>();
        SpiderStart bean = context.getBean(SpiderStart.class);

        Spider spider = bean.getSpider();


        long pageCount = spider.getPageCount();

        if (spider.getScheduler() instanceof MonitorableScheduler) {
            MonitorableScheduler scheduler = (MonitorableScheduler) spider.getScheduler();
            int leftRequestsCount = scheduler.getLeftRequestsCount(spider);
            int totalRequestsCount = scheduler.getTotalRequestsCount(spider);


            addMetric(metricSet, SPIDER_SUFFIX + "requestsCount", totalRequestsCount);

            addMetric(metricSet, SPIDER_SUFFIX + "requestsCount.left", leftRequestsCount);

        }

        SpiderMonitor.MonitorSpiderListener spiderListener = (SpiderMonitor.MonitorSpiderListener) spider.getSpiderListeners().get(0);
        AtomicInteger errorCount = spiderListener.getErrorCount();
        AtomicInteger successCount = spiderListener.getSuccessCount();

        addMetric(metricSet, SPIDER_SUFFIX + "page", pageCount);

        addMetric(metricSet, SPIDER_SUFFIX + "page.successCount", successCount);

        addMetric(metricSet, SPIDER_SUFFIX + "page.errorCount", errorCount);


        return metricSet;
    }

    private <T extends Number> void addMetric(Set<Metric<?>> metrics, String name,
                                              T value) {
        if (value != null) {
            metrics.add(new Metric<T>(name, value));
        }
    }
}
